from flask import Blueprint, request, jsonify, g
import sqlite3
import time
import json
import pandas as pd
import numpy as np
from database import get_db, DatabaseConnector
from utils.calculation_engine import calculate_sales_forecast, calculate_investment_return, calculate_market_analysis

calculation_bp = Blueprint('calculation', __name__)

@calculation_bp.route('/<int:scenario_id>/start', methods=['POST'])
def start_calculation(scenario_id):
    """开始测算"""
    data = request.json or {}
    calculation_mode = data.get('mode', 'quick')
    
    # 检查场景是否存在
    db = get_db()
    scenario = db.execute(
        'SELECT * FROM scenarios WHERE id = ?', 
        (scenario_id,)
    ).fetchone()
    
    if scenario is None:
        return jsonify({'error': '测算场景不存在'}), 404
    
    try:
        # 创建测算记录
        cursor = db.execute(
            '''INSERT INTO calculation_results 
               (scenario_id, result_data, status, started_at) 
               VALUES (?, ?, ?, CURRENT_TIMESTAMP)''',
            (scenario_id, '{}', 'running')
        )
        result_id = cursor.lastrowid
        db.commit()
        
        # 获取测算参数
        params_cursor = db.execute(
            'SELECT param_name, param_value, param_type FROM calculation_params WHERE scenario_id = ?',
            (scenario_id,)
        )
        params = {}
        for param in params_cursor.fetchall():
            # 根据参数类型转换值
            param_name = param['param_name']
            param_value = param['param_value']
            param_type = param['param_type']
            
            if param_type == 'number':
                try:
                    params[param_name] = float(param_value)
                except ValueError:
                    params[param_name] = 0
            elif param_type == 'boolean':
                params[param_name] = param_value.lower() == 'true'
            else:
                params[param_name] = param_value
        
        # 获取测算数据
        calculation_data = None
        
        # 检查是否有导入的数据
        import_record = db.execute(
            'SELECT file_path FROM data_imports WHERE scenario_id = ? ORDER BY imported_at DESC LIMIT 1',
            (scenario_id,)
        ).fetchone()
        
        if import_record:
            # 从导入的文件获取数据
            try:
                if import_record['file_path'].endswith('.csv'):
                    df = pd.read_csv(import_record['file_path'])
                elif import_record['file_path'].endswith(('.xlsx', '.xls')):
                    df = pd.read_excel(import_record['file_path'])
                elif import_record['file_path'].endswith('.json'):
                    df = pd.read_json(import_record['file_path'])
                
                calculation_data = df.to_dict('records')
            except Exception as e:
                # 更新测算状态为失败
                db.execute(
                    '''UPDATE calculation_results SET 
                       status = ?, result_data = ?, completed_at = CURRENT_TIMESTAMP
                       WHERE id = ?''',
                    ('failed', json.dumps({'error': f'读取数据失败: {str(e)}'}), result_id)
                )
                db.commit()
                return jsonify({'error': f'读取数据失败: {str(e)}'}), 500
        
        # 如果没有导入数据，检查是否关联了数据库
        elif scenario['db_connection_id']:
            # 从数据库获取数据
            conn_config = db.execute(
                'SELECT * FROM db_connections WHERE id = ?', 
                (scenario['db_connection_id'],)
            ).fetchone()
            
            if conn_config:
                # 构建查询获取相关数据
                # 根据场景类型使用不同的查询
                if scenario['scenario_type'] == 'sales':
                    query = """
                    SELECT date, product_category, product_name, region, quantity, amount 
                    FROM sales_data 
                    WHERE date >= ? AND date <= ?
                    """
                    # 使用参数中的时间范围
                    start_date = params.get('time_start', '2023-01-01')
                    end_date = params.get('time_end', '2023-12-31')
                    query_params = (start_date, end_date)
                else:
                    # 默认查询
                    query = "SELECT * FROM data LIMIT 1000"
                    query_params = ()
                
                # 执行查询
                connection_params = {
                    'db_type': conn_config['db_type'],
                    'host': conn_config['host'],
                    'port': conn_config['port'],
                    'database': conn_config['database'],
                    'username': conn_config['username'],
                    'password': conn_config['password']
                }
                
                connector = DatabaseConnector(connection_params)
                result, error = connector.execute_query(query, query_params)
                connector.close()
                
                if error:
                    # 更新测算状态为失败
                    db.execute(
                        '''UPDATE calculation_results SET 
                           status = ?, result_data = ?, completed_at = CURRENT_TIMESTAMP
                           WHERE id = ?''',
                        ('failed', json.dumps({'error': f'数据库查询失败: {error}'}), result_id)
                    )
                    db.commit()
                    return jsonify({'error': f'数据库查询失败: {error}'}), 500
                
                calculation_data = result
        
        # 如果没有数据，返回错误
        if calculation_data is None or len(calculation_data) == 0:
            # 更新测算状态为失败
            db.execute(
                '''UPDATE calculation_results SET 
                   status = ?, result_data = ?, completed_at = CURRENT_TIMESTAMP
                   WHERE id = ?''',
                ('failed', json.dumps({'error': '没有可用的测算数据'}), result_id)
            )
            db.commit()
            return jsonify({'error': '没有可用的测算数据'}), 400
        
        # 根据场景类型执行不同的测算
        start_time = time.time()
        calculation_result = None
        
        try:
            if scenario['scenario_type'] == 'sales':
                # 销售预测
                calculation_result = calculate_sales_forecast(
                    calculation_data, 
                    params,
                    detailed=calculation_mode == 'detailed'
                )
            elif scenario['scenario_type'] == 'investment':
                # 投资回报分析
                calculation_result = calculate_investment_return(
                    calculation_data, 
                    params,
                    detailed=calculation_mode == 'detailed'
                )
            elif scenario['scenario_type'] == 'market':
                # 市场分析
                calculation_result = calculate_market_analysis(
                    calculation_data, 
                    params,
                    detailed=calculation_mode == 'detailed'
                )
            else:
                # 默认测算逻辑
                calculation_result = {
                    'status': 'completed',
                    'summary': {
                        'total_items': len(calculation_data),
                        'calculation_time': time.time() - start_time,
                        'message': f'已完成{scenario["scenario_type"]}类型测算'
                    },
                    'data': calculation_data[:10]  # 只返回前10条数据作为示例
                }
        except Exception as e:
            # 测算过程中出错
            db.execute(
                '''UPDATE calculation_results SET 
                   status = ?, result_data = ?, completed_at = CURRENT_TIMESTAMP
                   WHERE id = ?''',
                ('failed', json.dumps({'error': f'测算过程出错: {str(e)}'}), result_id)
            )
            db.commit()
            return jsonify({'error': f'测算过程出错: {str(e)}'}), 500
        
        # 更新测算结果
        db.execute(
            '''UPDATE calculation_results SET 
               status = ?, result_data = ?, completed_at = CURRENT_TIMESTAMP
               WHERE id = ?''',
            ('completed', json.dumps(calculation_result), result_id)
        )
        db.commit()
        
        return jsonify({
            'result_id': result_id,
            'status': 'completed',
            'message': '测算已完成'
        })
        
    except sqlite3.Error as e:
        db.rollback()
        return jsonify({'error': f'测算启动失败: {str(e)}'}), 500

@calculation_bp.route('/<int:scenario_id>/status', methods=['GET'])
def get_calculation_status(scenario_id):
    """获取测算状态"""
    db = get_db()
    result = db.execute(
        '''SELECT * FROM calculation_results 
           WHERE scenario_id = ? 
           ORDER BY started_at DESC LIMIT 1''',
        (scenario_id,)
    ).fetchone()
    
    if result is None:
        return jsonify({'status': 'not_started', 'message': '尚未开始测算'})
    
    return jsonify({
        'result_id': result['id'],
        'status': result['status'],
        'started_at': result['started_at'],
        'completed_at': result['completed_at']
    })

@calculation_bp.route('/results/<int:result_id>', methods=['GET'])
def get_calculation_result(result_id):
    """获取测算结果"""
    db = get_db()
    result = db.execute(
        'SELECT * FROM calculation_results WHERE id = ?',
        (result_id,)
    ).fetchone()
    
    if result is None:
        return jsonify({'error': '测算结果不存在'}), 404
    
    # 获取关联的场景信息
    scenario = db.execute(
        'SELECT name, scenario_type FROM scenarios WHERE id = ?',
        (result['scenario_id'],)
    ).fetchone()
    
    return jsonify({
        'id': result['id'],
        'scenario_id': result['scenario_id'],
        'scenario_name': scenario['name'] if scenario else None,
        'scenario_type': scenario['scenario_type'] if scenario else None,
        'status': result['status'],
        'started_at': result['started_at'],
        'completed_at': result['completed_at'],
        'result_data': json.loads(result['result_data'])
    })

@calculation_bp.route('/<int:scenario_id>/cancel', methods=['POST'])
def cancel_calculation(scenario_id):
    """取消正在进行的测算"""
    db = get_db()
    # 查找最近的运行中的测算
    result = db.execute(
        '''SELECT * FROM calculation_results 
           WHERE scenario_id = ? AND status = 'running'
           ORDER BY started_at DESC LIMIT 1''',
        (scenario_id,)
    ).fetchone()
    
    if result is None:
        return jsonify({'error': '没有正在进行的测算'}), 404
    
    try:
        db.execute(
            '''UPDATE calculation_results SET 
               status = ?, completed_at = CURRENT_TIMESTAMP
               WHERE id = ?''',
            ('cancelled', result['id'])
        )
        db.commit()
        return jsonify({'message': '测算已取消'})
    except sqlite3.Error as e:
        db.rollback()
        return jsonify({'error': f'取消测算失败: {str(e)}'}), 500
